import sys
import random
from PyQt5.QtCore import Qt, QBasicTimer
from PyQt5.QtGui import QPainter, QColor
from PyQt5.QtWidgets import QApplication, QMainWindow, QFrame


class SnakeGame(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('贪吃蛇')
        self.setGeometry(100, 100, 640, 480)

        self.board = GameBoard(self)
        self.setCentralWidget(self.board)

        self.show()


class GameBoard(QFrame):
    SPEED = 100
    DOT_SIZE = 10
    MAX_X = 640 // DOT_SIZE
    MAX_Y = 480 // DOT_SIZE

    def __init__(self, parent):
        super().__init__(parent)
        self.timer = QBasicTimer()
        self.snake = [(10, 10), (10, 11), (10, 12)]
        self.food = self.random_food()
        self.direction = Qt.Key_Right
        self.in_game = True
        self.timer.start(self.SPEED, self)

    def random_food(self):
        while True:
            food = (random.randint(0, self.MAX_X - 1), random.randint(0, self.MAX_Y - 1))
            if food not in self.snake:
                return food

    def paintEvent(self, event):
        painter = QPainter(self)
        if self.in_game:
            self.draw_snake(painter)
            self.draw_food(painter)
        else:
            self.game_over(painter)

    def draw_snake(self, painter):
        painter.setBrush(QColor(0, 255, 0))
        for x, y in self.snake:
            painter.drawRect(x * self.DOT_SIZE, y * self.DOT_SIZE, self.DOT_SIZE, self.DOT_SIZE)

    def draw_food(self, painter):
        painter.setBrush(QColor(255, 0, 0))
        x, y = self.food
        painter.drawRect(x * self.DOT_SIZE, y * self.DOT_SIZE, self.DOT_SIZE, self.DOT_SIZE)

    def game_over(self, painter):
        painter.setBrush(QColor(0, 0, 0))
        painter.drawText(self.rect(), Qt.AlignCenter, "游戏结束")

    def keyPressEvent(self, event):
        if not self.in_game:
            return
        key = event.key()
        if key in [Qt.Key_Left, Qt.Key_Right, Qt.Key_Up, Qt.Key_Down]:
            if (key == Qt.Key_Left and self.direction != Qt.Key_Right) or \
               (key == Qt.Key_Right and self.direction != Qt.Key_Left) or \
               (key == Qt.Key_Up and self.direction != Qt.Key_Down) or \
               (key == Qt.Key_Down and self.direction != Qt.Key_Up):
                self.direction = key

    def timerEvent(self, event):
        if event.timerId() == self.timer.timerId():
            self.move_snake()
        else:
            super(GameBoard, self).timerEvent(event)

    def move_snake(self):
        head_x, head_y = self.snake[0]
        if self.direction == Qt.Key_Left:
            head_x -= 1
        elif self.direction == Qt.Key_Right:
            head_x += 1
        elif self.direction == Qt.Key_Up:
            head_y -= 1
        elif self.direction == Qt.Key_Down:
            head_y += 1

        new_head = (head_x, head_y)

        if new_head in self.snake or \
           head_x < 0 or head_x >= self.MAX_X or \
           head_y < 0 or head_y >= self.MAX_Y:
            self.in_game = False
            self.timer.stop()
            self.update()
            return

        self.snake = [new_head] + self.snake[:-1]

        if new_head == self.food:
            self.snake.append(self.snake[-1])
            self.food = self.random_food()

        self.update()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    game = SnakeGame()
    sys.exit(app.exec_())
